"""
读取THUC NEWS数据集信息
"""
import os
import sys
import zipfile

import pandas as pd
from tqdm import tqdm

from config.sys_config import THUC_NEWS_DATASET_PATH
from dataset._base_uitls import download_dataset, support_filename_encode

# 数据集下载地址
THUC_NEWS_DATASET_URL = "https://thunlp.oss-cn-qingdao.aliyuncs.com/THUCNews.zip"


def read_thuc_news(download=False, chunk_size=1024 * 10):
    if not os.path.exists(THUC_NEWS_DATASET_PATH):
        if download:
            download_dataset(THUC_NEWS_DATASET_PATH, THUC_NEWS_DATASET_URL, chunk_size)
        else:
            print("Error: Dataset is not exist.Please use arg `download=True` to start download.")
            return sys.exit(1)
    return support_filename_encode(zipfile.ZipFile(THUC_NEWS_DATASET_PATH))


class ThucNewsDataset:
    _THUC_NEWS_FILE: zipfile.ZipFile
    _FILENAME_DATA_FRAME: pd.DataFrame
    _TRAIN_DATA_FRAME: pd.DataFrame
    _TEST_DATA_FRAME: pd.DataFrame
    seedNum: int
    train_scale: float

    def __init__(self, seedNum=66, train_scale=0.8,
                 zipFile: zipfile.ZipFile = None, download=True, chunk_size=1024 * 10):
        self.seedNum = seedNum
        self.train_scale = train_scale

        if zipFile is None:
            self._THUC_NEWS_FILE = read_thuc_news(download, chunk_size)
        else:
            self._THUC_NEWS_FILE = zipFile
        self._read_file_list()
        self._create_data_frame()

    def _create_data_frame(self):
        self._TRAIN_DATA_FRAME = self._FILENAME_DATA_FRAME.groupby('label', group_keys=False)\
            .apply(lambda x: x.sample(frac=self.train_scale)).copy()
        self._TEST_DATA_FRAME = self._FILENAME_DATA_FRAME[~self._FILENAME_DATA_FRAME.index.isin(
            self._TRAIN_DATA_FRAME.index)].copy()

    def _read_file_list(self):
        _filename_list = []
        for _filename in tqdm(self._THUC_NEWS_FILE.namelist(), desc="正在读取文件信息"):
            if _filename.endswith('.txt') and not _filename.startswith('__MACOSX'):
                content = str(self._THUC_NEWS_FILE.read(_filename), encoding='utf-8')
                path, name = _filename.rsplit("/", 1)
                label = path.split('/')[-1]
                file_id = name.replace('.txt', '')
                row = (label, int(file_id), path, name, content)
                _filename_list.append(row)
        self._FILENAME_DATA_FRAME = pd.DataFrame(_filename_list,
                                                 columns=['label', 'file_id', 'path', 'filename', 'content'])
        self._THUC_NEWS_FILE.close()

    @property
    def file_data(self) -> pd.DataFrame:
        return self._FILENAME_DATA_FRAME.copy()

    @property
    def train_data(self) -> pd.DataFrame:
        return self._TRAIN_DATA_FRAME.copy()

    @property
    def test_data(self) -> pd.DataFrame:
        return self._TEST_DATA_FRAME.copy()


if __name__ == '__main__':
    thuc_news_file = read_thuc_news(download=True)
    thuc_news_dataset = ThucNewsDataset(zipFile=thuc_news_file)
    print(thuc_news_dataset.file_data.size)
    print(thuc_news_dataset.train_data.index)
    print(thuc_news_dataset.test_data.index)
